
#include <vector>
#include "AESSymmetricKey.h"
#include "ECDSAPublicKey.h"
#include "DSAPublicKey.h"
#include "RSAPublicKey.h"
#include "RSAPrivateKey.h"
#include "CryptEChain.h"
#include "DecryptWrapper.h"

/** General initialisation function for the decryption wrapper.
 * 
 * @param deckey
 * DecryptionKey to use for decrypting data if present and not NULL.
 * 
 * @param chkkey
 * CheckingKey to use to verify signatures on the data stream if
 * present and not NULL.
 * 
 * @param decompress
 * A flag to specify an additional decompression step if true. */
void DecryptWrapper::initDecryptWrapper(DecryptionKey *deckey, CheckingKey *chkkey, int decompress)
{
   if (deckey)
   {
      switch (deckey->getKeyType())
      {
         case Key::KEY_TYPE_PRIVATE_RSA:
            { deckeycopy = new RSAPrivateKey(*((RSAPrivateKey *)deckey)); break; }
         case Key::KEY_TYPE_SYMM_AES:
            { deckeycopy = new AESSymmetricKey(*((AESSymmetricKey *)deckey)); break; }
         default:
            { errprint("decrypt wrapper has no key: can't do decryption!"); deckeycopy = NULL; break; }
      }
   } else
   {
      deckeycopy = NULL;
   }
   if (chkkey)
   {
      switch (chkkey->getKeyType())
      {
         case Key::KEY_TYPE_PUBLIC_RSA:
            { chkkeycopy = new RSAPublicKey(*((RSAPublicKey *)deckey)); break; }
         case Key::KEY_TYPE_PUBLIC_DSA:
            { chkkeycopy = new DSAPublicKey(*((DSAPublicKey *)deckey)); break; }
         case Key::KEY_TYPE_PUBLIC_ECDSA:
            { chkkeycopy = new ECDSAPublicKey(*((ECDSAPublicKey *)deckey)); break; }
         default:
            { errprint("unknown signature check key type"); chkkeycopy = NULL; break; }
      }
   } else
   {
      chkkeycopy = NULL;
   }
   dodecompress = (decompress != 0);
   wastrusted = 0;
   apr_status_t status;
   status = apr_pool_create(&mp, NULL);
   CheckAPRError(status);
}

/** Construct a decryption process involving just decryption with the given
 * DecryptionKey.
 * 
 * @param deckey
 * DecryptionKey to use for decrypting data if present and not NULL. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey)
{
   initDecryptWrapper(deckey, NULL, FALSE);
}

/** Construct a decryption process which verifies that the signature on the
 * data matches the given CheckingKey, before decrypting with the given
 * DecryptionKey.
 * 
 * @param deckey
 * DecryptionKey to use for decrypting data if present and not NULL.
 * 
 * @param chkkey
 * CheckingKey to use to verify signatures on the data stream if
 * present and not NULL. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey, CheckingKey *chkkey)
{
   initDecryptWrapper(deckey, chkkey, FALSE);
}

/** Construct a decryption process which decrypts with the given DecryptionKey,
 * but also optionally apply a decompression step to the resulting data.
 * 
 * @param deckey
 * DecryptionKey to use for decrypting data if present and not NULL.
 * 
 * @param dodecompress
 * A flag to specify an additional decompression step if true. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey, int dodecompress)
{
   initDecryptWrapper(deckey, NULL, dodecompress);
}

/** Construct a decryption process which verifies that the signature on the
 * data matches the given CheckingKey, decrypts with the given DecryptionKey
 * and can also also optionally apply a decompression step to the decrypted 
 * data.
 * 
 * @param deckey
 * DecryptionKey to use for decrypting data if present and not NULL.
 * 
 * @param chkkey
 * CheckingKey to use to verify signatures on the data stream if
 * present and not NULL.
 * 
 * @param dodecompress
 * A flag to specify an additional decompression step if true. */
DecryptWrapper::DecryptWrapper(DecryptionKey *deckey, CheckingKey *chkkey, int dodecompress)
{
   initDecryptWrapper(deckey, chkkey, dodecompress);
}

/** Destructor for the decryption wrapper. */
DecryptWrapper::~DecryptWrapper()
{
   if (deckeycopy) delete deckeycopy;
   if (chkkeycopy) delete chkkeycopy;
   apr_pool_destroy(mp);
}

/** Apply the decryption process to a buffer of data, returning a vector of
 * decoded data.
 * 
 * @param buffer
 * Buffer for encoded data input.
 * 
 * @param length
 * Length of buffer for data input.
 * 
 * @return
 * An std::vector representing a variable-sized buffer into which the decoded
 * data is put. The caller is responsible for deleting the vector instance. */
std::vector<char> *DecryptWrapper::decryptBuffer(const char *buffer, size_t length)
{
   int smallblock = (ENV_READ_BUFFER > (length/2)) ? (length/2) : ENV_READ_BUFFER;
   size_t nlength = length;
   APRBucketBrigade *inputbb = new APRBucketBrigade(mp);
   if (inputbb->bufferToBucketBrigade(buffer, &length, smallblock) != 0)
   {
      errprint("failed to push the buffer into a bucket");
      delete inputbb;
      return NULL;
   }
   if (nlength != length)
   {
      errprint("bucket generation failed, pushed %d bytes to bucket, wanted %d", nlength, length);
      delete inputbb;
      return NULL;
   }
   APRBucketBrigade *outputbb = decryptToBucket(inputbb);
   delete inputbb;
   char *outputbuffer = NULL;
   size_t ilength = 0;
   std::vector<char> *newstring = NULL;
   if (outputbb->bucketBrigadeToBuffer(&outputbuffer, &ilength) == 0)
   {
      // outputbuffer cannot be NULL!
      newstring = new std::vector<char>();
      newstring->insert(newstring->end(), outputbuffer, &(outputbuffer[ilength]));
   } else
   {
      errprint("failed to put bucket into buffer");
   }
   if (outputbuffer) free(outputbuffer);
   delete outputbb;
   return newstring;
}

/** Apply the decryption process to a pipe of data in the form of an
 * APRBucketBrigade, returning a pointer to a newly created APRBucketBrigade
 * containing the decoded data.
 * 
 * @param bbrgde
 * The bucket brigade input of encoded data.
 * 
 * @return
 * New bucket brigade representing the complete decoded data stream. The
 * caller is responsible for deleting the APRBucketBrigade instance. */
APRBucketBrigade *DecryptWrapper::decryptToBucket(APRBucketBrigade *bbrgde)
{
   CryptEChain *crypty = new CryptEChain();
   crypty->buildDecryptionChain2(deckeycopy, chkkeycopy, dodecompress, bbrgde->realLength());
   APRBucketBrigade *outputbrgde = new APRBucketBrigade(mp);
   int prociter = 0;
   while (prociter == 0)
      prociter = crypty->processIteration(bbrgde, outputbrgde);
   delete bbrgde;
   wastrusted = crypty->isTrusted();
   delete crypty;
   if (prociter == -1)
   {
      delete outputbrgde;
      return NULL;
   }
   return outputbrgde;
}
